"""
Chat agent implementation for basic conversation
"""
from typing import Dict, List, Optional, Any
import json
import re
from openai import OpenAI
from .base import BaseAgent
from ..core.message import Message, MessageType
from ..core.prompt import wrap_content, unwrap_message


class ChatAgent(BaseAgent):
    """聊天智能体类，用于基本对话"""

    def __init__(self,
                 agent_id: str,
                 api_key: str,
                 model: str = "gpt-3.5-turbo",
                 system_prompt: str = "你是一个友好的AI助手。",
                 duty: str = "协助用户完成任务。",
                 temperature: float = 0.7,
                 base_url: Optional[str] = None,
                 max_num_context: int = 1000,
                 use_protocol: bool = False):
        """
        初始化聊天智能体
        Args:
            agent_id: 智能体ID
            api_key: OpenAI API密钥
            model: 使用的模型名称
            system_prompt: 系统提示词
            temperature: 模型温度参数
            base_url: API基础URL，用于支持其他兼容OpenAI接口的模型服务
        """
        super().__init__(agent_id, max_num_context, use_protocol, system_prompt, duty)
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url if base_url else "https://api.openai.com/v1")
        self.model = model
        self.temperature = temperature

    def forward(self, message: Message) -> Message:
        """
        处理接收到的消息
        Args:
            message: 输入消息
        Returns:
            响应消息
        """
        # 将消息添加到会话上下文
        self.add_context(content=json.dumps({
            "role":
            "user" if message.sender != self.agent_id else "assistant",
            "content":
            message.content
        }))

        # 调用OpenAI获取响应
        response = self._call_openai()
        content = self._clean_think(response.choices[0].message.content) # 清除大模型深度思考内容
        self.add_context(
            content=json.dumps({
                "role": "assistant",
                "content": content
            }),
            # content=response.choices[0].message.model_dump_json(),
        )

        # 创建响应消息
        msg = self.create_message(
            content=content,
            msg_type=MessageType.TEXT,
            # receiver=message.sender
            metadata=message.metadata)
        
        if self.use_protocol:
            try:
                msg.receiver = unwrap_message(msg.content)['to']
            except Exception as e:
                print(f'无法解析 Sonetto 消息: {msg.content}')
                raise e

        return msg

    def _clean_think(self, content):
        cleaned_text = re.sub(r'<think[^>]*>.*?</think>',
                              '',
                              content,
                              flags=re.DOTALL)
        return re.sub(r'^\n+', '', cleaned_text)

    def _call_openai(self) -> Any:
        """调用OpenAI API"""
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=self._build_messages(),
                temperature=self.temperature)

            # 检查响应是否遵循协议
            if self.use_protocol:
                while unwrap_message(response.choices[0].message.content) is None:
                    print(f'无法解析 Sonetto 消息！！！')
                    self.add_context(json.dumps({
                        'role': 'assistant',
                        'content': response.choices[0].message.content
                    }))
                    self.add_context(json.dumps({'role': 'user', 'content': 'XML解析器无法解析到sonetto-message元素。\
                                      请把你发的上一条消息用Sonetto通信协议包装并重新发送。请不要写错你的消息传递对象。'}))
                    response = self.client.chat.completions.create(
                        model=self.model,
                        messages=self._build_messages(),
                        temperature=self.temperature)
                    return response
            # print(response.choices[0].message.content)
            return response
        except Exception as e:
            raise e
            # return f"Error calling OpenAI API: {str(e)}"

    def _build_messages(self) -> List[Dict]:
        """构建OpenAI API所需的消息历史"""
        messages = []

        # 添加系统提示
        if self.sonetto_prompt:
            messages.append({"role": "system", "content": self.sonetto_prompt})
        messages.append({"role": "system", "content": self.system_prompt})

        # 添加会话历史
        for ctx in self.get_context():
            try:
                msg = json.loads(ctx["content"])
                messages.append({
                    "role": msg["role"],
                    "content": msg["content"]
                })
            except:
                continue

        return messages

    def set_system_prompt(self, prompt: str):
        """
        设置系统提示词
        Args:
            prompt: 新的系统提示词
        """
        self.system_prompt = prompt

    def set_temperature(self, temperature: float):
        """
        设置模型温度参数
        Args:
            temperature: 新的温度值 (0.0-2.0)
        """
        if not 0 <= temperature <= 2:
            raise ValueError("Temperature must be between 0 and 2")
        self.temperature = temperature
